import argparse
import os
import shutil

import yaml
from imutils import paths
from sklearn.model_selection import train_test_split
from tqdm import tqdm


def createDatasets(datasets, dirname):
    dataDir = os.path.join(data_dict[dirname], name)
    if not os.path.exists(dataDir):
        os.makedirs(dataDir)
    for datapath in datasets:
        shutil.copy(datapath, dataDir)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="获取杭州人工认为有缺陷大图")
    parser.add_argument('--data',
                        default=r"../data/abnormal.yaml",
                        help="没有分的文件夹")
    opt = parser.parse_args()

    with open(opt.data,encoding="utf-8") as f:
        data_dict = yaml.load(f, Loader=yaml.FullLoader)  # data dict

    classesData = {}
    for name in tqdm(data_dict["names"]):
        nameImgs = list(paths.list_images(os.path.join(data_dict["allDatas"], name)))
        X_train, X_test_val, _, _ = train_test_split(nameImgs, nameImgs, test_size=0.4, random_state=1024)
        X_test, X_val, _, _ = train_test_split(X_test_val, X_test_val, test_size=0.5, random_state=1024)
        createDatasets(X_train, "train")
        createDatasets(X_test, "test")
        createDatasets(X_val, "val")
